package com.dynamic.datasource.router;

import com.dynamic.datasource.connection.ConnectionFactory;
import com.dynamic.datasource.connection.ConnectionProxy;
import com.dynamic.datasource.exception.NoSuchDataSourceException;
import com.dynamic.datasource.transaction.TransactionContext;
import org.springframework.jdbc.datasource.AbstractDataSource;
import org.springframework.util.StringUtils;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;

/**
 * @author 王先堂
 * @Description 抽象数据源路由
 * @createTime 2022/03/24/ 15:27:00
 */
public abstract class AbstractDataSourceRouter extends AbstractDataSource {

    /**
     * 决定数据源
     *
     * @return 目标数据源
     */
    protected abstract DataSource determineDataSource();

    @Override
    public Connection getConnection() throws SQLException {
        String id = TransactionContext.get();
        if (StringUtils.isEmpty(id)) {
            return determineDataSource().getConnection();
        }
        String name = DataSourceContext.get();
        if (StringUtils.isEmpty(name)) {
            throw new NoSuchDataSourceException("no such DataSource " + name);
        }
        ConnectionProxy connection = ConnectionFactory.getConnection(name);
        return connection == null ? getConnectionProxy(name, determineDataSource().getConnection()) : connection;
    }

    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        String id = TransactionContext.get();
        if (StringUtils.isEmpty(id)) {
            return determineDataSource().getConnection(username, password);
        } else {
            String name = DataSourceContext.get();
            if (StringUtils.isEmpty(name)) {
                throw new NoSuchDataSourceException("no such DataSource " + name);
            }
            ConnectionProxy connection = ConnectionFactory.getConnection(name);
            return connection == null ? getConnectionProxy(name, determineDataSource().getConnection(username, password))
                    : connection;
        }
    }

    /**
     * 获取代理连接对象
     *
     * @param name       数据源名称
     * @param connection 连接对象
     * @return
     */
    private Connection getConnectionProxy(String name, Connection connection) {
        ConnectionProxy connectionProxy = new ConnectionProxy(connection, name);
        ConnectionFactory.putConnection(name, connectionProxy);
        return connectionProxy;
    }

    @Override
    public <T> T unwrap(Class<T> iface) throws SQLException {
        if (iface.isInstance(this)) {
            return (T) this;
        }
        return determineDataSource().unwrap(iface);
    }

    @Override
    public boolean isWrapperFor(Class<?> iface) throws SQLException {
        return (iface.isInstance(this) || determineDataSource().isWrapperFor(iface));
    }
}
